import numpy as np
from zipline.api import order, symbol, record
from zipline import run_algorithm
from datetime import datetime
import pytz

def initialize(context):
    # 设置策略参数
    context.lookback = 30
    context.zscore_threshold = 2.0
    context.asset = symbol('AAPL')  # 示例股票，可以根据需要修改

def handle_data(context, data):
    # 获取历史价格数据
    prices = data.history(context.asset, 'close', context.lookback, '1d')
    
    # 计算Z-Score
    mean = np.mean(prices)
    std = np.std(prices)
    zscore = (prices[-1] - mean) / std
    
    # 获取当前持仓
    current_position = context.portfolio.positions[context.asset].amount
    
    # 交易逻辑
    if zscore < -context.zscore_threshold and current_position <= 0:
        # 买入
        order(context.asset, 10)  # 买入10股
    elif zscore > context.zscore_threshold and current_position > 0:
        # 卖出
        order(context.asset, -10)  # 卖出10股

# 以下代码用于本地测试
if __name__ == '__main__':
    start = datetime(2011, 1, 1, tzinfo=pytz.UTC)
    end = datetime(2012, 1, 1, tzinfo=pytz.UTC)
    capital_base = 10000
    result = run_algorithm(start=start, end=end, initialize=initialize, 
                          capital_base=capital_base, handle_data=handle_data, 
                          bundle='quandl')